package com.shanzmoo.core.util;

import cn.hutool.core.io.FileUtil;
import com.hankcs.hanlp.dictionary.CustomDictionary;
import com.hankcs.hanlp.seg.common.Term;
import com.hankcs.hanlp.tokenizer.StandardTokenizer;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * Hanlp 分词工具
 *
 * @author: Gzhao 2020/12/24
 */
@Slf4j
public class HanlpUtils {

    private static final List<String> NATURE_LIST = Arrays.asList("mq", "vn");

    static {
        List<String> myWords = readMyDictionary();

        // 动态增加
        myWords.forEach(e->{
            CustomDictionary.add(e);
        });
    }

    private HanlpUtils() {
    }

    public static List<String> segment(String text) {
        log.info("===== Hanlp 分词开始 =====");
        log.info("原始数据：{}", text);
        List<Term> termList = StandardTokenizer.segment(text);

        List<String> result = termList.stream()
                .filter(e -> NATURE_LIST.contains(e.nature.toString()) || e.nature.startsWith('n') || e.nature.startsWith('g'))
                .map(e -> {
                    return e.word;
                })
                .distinct()
                .collect(Collectors.toList());

        log.info("分词后数据：{}", result);
        log.info("===== Hanlp 分词结束 =====");
        return result;
    }

    private static List<String> readMyDictionary() {
        try {
            File file = FileUtil.touch("myDictionary.txt");

            byte[] bytes = FileUtil.readBytes(file);

            String content = new String(bytes);
            log.info("加载自定义分词字典：\n{}", content);

            String[] wordList = content.split("\r\n");
            return Arrays.asList(wordList);

        } catch (Exception e) {
            log.error("Hanlp加载自定义分词失败，未找到自定义词典");
            return Collections.emptyList();
        }
    }
}
